{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5eebfc21-b54a-435d-84dc-b9de1d620e4e",
   "metadata": {},
   "source": [
    "# Tool Calling\n",
    "\n",
    "In this example we'll go over a basic application that selects between a few different tools, calls them, and uses an LLM to process the results.\n",
    "\n",
    "Requirements:\n",
    "- OpenAI API key, set as the env variable `OPENAI_API_KEY`\n",
    "- [tomorrow.io](tomorrow.io) API Key, set as the env variable `TOMORROW_API_KEY`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "537a0463-1d35-48dd-b0af-9ebabf2b08dc",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b60de9a5-43b9-421d-b2c4-a758d3fcde3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "import json\n",
    "import os\n",
    "from typing import Callable, Optional\n",
    "\n",
    "import openai\n",
    "import requests\n",
    "\n",
    "from burr.core import State, action, when\n",
    "from burr.core.application import ApplicationBuilder\n",
    "\n",
    "# Set up + instantiate instrumentation\n",
    "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n",
    "\n",
    "OpenAIInstrumentor().instrument()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cc27a92-e36f-4306-91b8-aa68f7d2e6a2",
   "metadata": {},
   "source": [
    "# Defining tools\n",
    "\n",
    "Let's define a few tools that we want to use. They'll all have:\n",
    "\n",
    "1. Primitive input, annotated with types\n",
    "2. return types of dictionaries (free-form)\n",
    "\n",
    "This way we can have some structure in the return type. You'll likely want a stricter return type for your cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "690935f7-19b4-494d-abab-9a3f7cf0c26c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _weather_tool(latitude: float, longitude: float) -> dict:\n",
    "    \"\"\"Queries the weather for a given latitude and longitude.\"\"\"\n",
    "    api_key = os.environ.get(\"TOMORROW_API_KEY\")\n",
    "    url = f\"https://api.tomorrow.io/v4/weather/forecast?location={latitude},{longitude}&apikey={api_key}\"\n",
    "\n",
    "    response = requests.get(url)\n",
    "    if response.status_code == 200:\n",
    "        return response.json()\n",
    "    else:\n",
    "        return {\"error\": f\"Failed to get weather data. Status code: {response.status_code}\"}\n",
    "\n",
    "\n",
    "def _text_wife_tool(message: str) -> dict:\n",
    "    \"\"\"Texts your wife with a message.\"\"\"\n",
    "    # Dummy implementation for the text wife tool\n",
    "    # Replace this with actual SMS API logic\n",
    "    return {\"action\": f\"Texted wife: {message}\"}\n",
    "\n",
    "\n",
    "def _order_coffee_tool(\n",
    "    size: str, coffee_preparation: str, any_modifications: Optional[str] = None\n",
    ") -> dict:\n",
    "    \"\"\"Orders a coffee with the given size, preparation, and any modifications.\"\"\"\n",
    "    # Dummy implementation for the order coffee tool\n",
    "    # Replace this with actual coffee shop API logic\n",
    "    return {\n",
    "        \"action\": (\n",
    "            f\"Ordered a {size} {coffee_preparation}\" + \"with {any_modifications}\"\n",
    "            if any_modifications\n",
    "            else \"\"\n",
    "        )\n",
    "    }\n",
    "\n",
    "\n",
    "def _fallback(response: str) -> dict:\n",
    "    \"\"\"Tells the user that the assistant can't do that -- this should be a fallback\"\"\"\n",
    "    return {\"response\": response}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c198e8d-a467-433c-92f0-b125d3529750",
   "metadata": {},
   "source": [
    "# Defining the tools (as OpenAI wants them)\n",
    "\n",
    "We do a little cleverness to get types -- you'll likely want to use pydantic to get this right later on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "884bf319-5648-4049-a654-ef2305f32ec0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You'll want to add more types here as needed\n",
    "TYPE_MAP = {\n",
    "    str: \"string\",\n",
    "    int: \"integer\",\n",
    "    float: \"number\",\n",
    "    bool: \"boolean\",\n",
    "}\n",
    "\n",
    "# You can also consider using a library like pydantic to further integrate with OpenAI\n",
    "OPENAI_TOOLS = [\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": fn_name,\n",
    "            \"description\": fn.__doc__ or fn_name,\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    param.name: {\n",
    "                        \"type\": TYPE_MAP.get(param.annotation)\n",
    "                        or \"string\",  # TODO -- add error cases\n",
    "                        \"description\": param.name,\n",
    "                    }\n",
    "                    for param in inspect.signature(fn).parameters.values()\n",
    "                },\n",
    "                \"required\": [param.name for param in inspect.signature(fn).parameters.values()],\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "    for fn_name, fn in {\n",
    "        \"query_weather\": _weather_tool,\n",
    "        \"order_coffee\": _order_coffee_tool,\n",
    "        \"text_wife\": _text_wife_tool,\n",
    "        \"fallback\": _fallback,\n",
    "    }.items()\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c27937-e723-4f4f-ad5f-70bc518919ba",
   "metadata": {},
   "source": [
    "# Defining our actions\n",
    "\n",
    "We're going to have:\n",
    "1. An action that processes the user input\n",
    "2. An action that selects the tool + parameters\n",
    "3. An action that calls the tool\n",
    "4. An action that formats the results\n",
    "\n",
    "This way each is decoupled. The action that calls the tool will be \"generic\", and we'll be binding it to each tool (see next section)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c45bea6d-f5af-4144-8d77-3d631d127b6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@action(reads=[], writes=[\"query\"])\n",
    "def process_input(state: State, query: str) -> State:\n",
    "    \"\"\"Simple action to process input for the assistant.\"\"\"\n",
    "    return state.update(query=query)\n",
    "\n",
    "\n",
    "\n",
    "@action(reads=[\"query\"], writes=[\"tool_parameters\", \"tool\"])\n",
    "def select_tool(state: State) -> State:\n",
    "    \"\"\"Selects the tool + assigns the parameters. Uses the tool-calling API.\"\"\"\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": (\n",
    "                \"You are a helpful assistant. Use the supplied tools to assist the user, if they apply in any way. Remember to use the tools! They can do stuff you can't.\"\n",
    "                \"If you can't use only the tools provided to answer the question but know the answer, please provide the answer\"\n",
    "                \"If you cannot use the tools provided to answer the question, use the fallback tool and provide a reason. \"\n",
    "                \"Again, if you can't use one tool provided to answer the question, use the fallback tool and provide a reason. \"\n",
    "                \"You must select exactly one tool no matter what, filling in every parameters with your best guess. Do not skip out on parameters!\"\n",
    "            ),\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": state[\"query\"]},\n",
    "    ]\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=messages,\n",
    "        tools=OPENAI_TOOLS,\n",
    "    )\n",
    "\n",
    "    # Extract the tool name and parameters from OpenAI's response\n",
    "    if response.choices[0].message.tool_calls is None:\n",
    "        return state.update(\n",
    "            tool=\"fallback\",\n",
    "            tool_parameters={\n",
    "                \"response\": \"No tool was selected, instead response was: {response.choices[0].message}.\"\n",
    "            },\n",
    "        )\n",
    "    fn = response.choices[0].message.tool_calls[0].function\n",
    "\n",
    "    return state.update(tool=fn.name, tool_parameters=json.loads(fn.arguments))\n",
    "\n",
    "\n",
    "@action(reads=[\"tool_parameters\"], writes=[\"raw_response\"])\n",
    "def call_tool(state: State, tool_function: Callable) -> State:\n",
    "    \"\"\"Action to call the tool. This will be bound to the tool function.\"\"\"\n",
    "    response = tool_function(**state[\"tool_parameters\"])\n",
    "    return state.update(raw_response=response)\n",
    "\n",
    "\n",
    "@action(reads=[\"query\", \"raw_response\"], writes=[\"final_output\"])\n",
    "def format_results(state: State) -> State:\n",
    "    \"\"\"Action to format the results in a usable way. Note we're not cascading in context for the chat history.\n",
    "    This is largely due to keeping it simple, but you'll likely want to pass IDs around or maintain the chat history yourself\n",
    "    \"\"\"\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": (\n",
    "                    \"You are a helpful assistant. Your goal is to take the\"\n",
    "                    \"data presented and use it to answer the original question:\"\n",
    "                ),\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": (\n",
    "                    f\"The original question was: {state['query']}.\"\n",
    "                    f\"The data is: {state['raw_response']}. Please format\"\n",
    "                    \"the data and provide a response that responds to the original query.\"\n",
    "                    \"As always, be concise (tokens aren't free!).\"\n",
    "                ),\n",
    "            },\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    return state.update(final_output=response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fccd23db-b8d6-4179-a452-b37defbeff51",
   "metadata": {},
   "source": [
    "# Building the application\n",
    "\n",
    "Now let's build the application. We will:\n",
    "1. Add the right actions\n",
    "2. Add the right transitions\n",
    "3. Set up a tracker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9477c4f1-96ad-4125-a504-ab47ac4e56d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_actions(\n",
    "        process_input,\n",
    "        select_tool,\n",
    "        format_results,\n",
    "        query_weather=call_tool.bind(tool_function=_weather_tool),\n",
    "        text_wife=call_tool.bind(tool_function=_text_wife_tool),\n",
    "        order_coffee=call_tool.bind(tool_function=_order_coffee_tool),\n",
    "        fallback=call_tool.bind(tool_function=_fallback),\n",
    "    )\n",
    "    .with_transitions(\n",
    "        (\"process_input\", \"select_tool\"),\n",
    "        (\"select_tool\", \"query_weather\", when(tool=\"query_weather\")),\n",
    "        (\"select_tool\", \"text_wife\", when(tool=\"text_wife\")),\n",
    "        (\"select_tool\", \"order_coffee\", when(tool=\"order_coffee\")),\n",
    "        (\"select_tool\", \"fallback\", when(tool=\"fallback\")),\n",
    "        ([\"query_weather\", \"text_wife\", \"order_coffee\", \"fallback\"], \"format_results\"),\n",
    "        (\"format_results\", \"process_input\"),\n",
    "    )\n",
    "    .with_entrypoint(\"process_input\")\n",
    "    .with_tracker(project=\"test_tool_calling\", use_otel_tracing=True)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b7d206e-0c4f-40d2-a4f1-21341f8d26b6",
   "metadata": {},
   "source": [
    "# Running the app\n",
    "\n",
    "And it works! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8b94805e-6da2-4fbe-aa05-92acd1e63ca4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Based on the data provided for San Francisco as of September 30, 2024, and subsequent days:\n",
      "\n",
      "- **Temperature**: Currently around 25.5°C to 25.9°C.\n",
      "- **Weather**: Clear skies, with cloud cover at 0%.\n",
      "- **Humidity**: Ranges between 35% to 46%.\n",
      "- **Wind**: Light winds from the north, shifting to the west, with speeds ranging from 1.63 to 3.48 km/h. Gusts can reach up to 4.17 km/h.\n",
      "- **Visibility**: Excellent, at 16 km.\n",
      "- **Pressure**: Around 1007.7 hPa.\n",
      "- **UV Index**: Moderate to high, peaking at 6.\n",
      "\n",
      "This indicates a clear, warm day with comfortable humidity levels and light winds.\n"
     ]
    }
   ],
   "source": [
    "app.visualize(output_file_path=\"./statemachine.png\")\n",
    "action, result, state = app.run(\n",
    "    halt_after=[\"format_results\"],\n",
    "    inputs={\"query\": \"What's the weather like in San Francisco?\"},\n",
    ")\n",
    "print(state[\"final_output\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd763c1-128b-4b78-a973-ac01723ec636",
   "metadata": {},
   "source": [
    "# Visualizing\n",
    "\n",
    "We can view in the UI. Make sure `burr` is runnng for this link to work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6764124d-eb0f-43f7-92e7-808893b9a831",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "url = f\"[Link to UI](http://localhost:7241/project/demo_email_assistant/{app.uid})\"\n",
    "Markdown(url)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
